"""
Alibaba Cloud remote deployment and benchmark execution
"""

import os
import time
from typing import Dict, Any, List
from fabric import Connection
from benchmark.utils import BenchError, Print, run_command, get_timestamp, ensure_directory
from alibaba.instance import AlibabaInstanceManager


class AlibabaBench:
    """Alibaba Cloud benchmark implementation for Dumbo"""
    
    def __init__(self, ctx):
        self.ctx = ctx
        self.root_dir = os.getcwd()
        self.data_dir = os.path.join(self.root_dir, "data")
        self.instance_manager = AlibabaInstanceManager()
        self.connections = []
        ensure_directory(self.data_dir)
    
    def run(self, bench_params: Dict[str, Any], node_params: Dict[str, Any]):
        """Run Alibaba Cloud benchmark"""
        Print.subheading("Starting Alibaba Cloud Dumbo benchmark")
        
        try:
            # Create instances
            nodes = bench_params['nodes'][0]
            instance_type = bench_params.get('instance_type', 'ecs.g6e.xlarge')
            region = bench_params.get('region', 'cn-hangzhou')
            
            instances = self.instance_manager.create_instances(nodes, instance_type, region)
            
            # Deploy to instances
            self.deploy(bench_params)
            
            # Run benchmark
            self._run_benchmark(bench_params, node_params)
            
            # Collect results
            self._collect_results()
            
            Print.success("Alibaba Cloud benchmark completed successfully")
            
        except Exception as e:
            raise BenchError(f"Alibaba Cloud benchmark failed: {e}")
        finally:
            # Cleanup
            self._cleanup()
    
    def deploy(self, bench_params: Dict[str, Any]):
        """Deploy Dumbo to Alibaba Cloud instances"""
        Print.subheading("Deploying Dumbo to Alibaba Cloud instances")
        
        nodes = bench_params['nodes'][0]
        instances = self.instance_manager.get_instance_info()
        
        if len(instances) < nodes:
            raise BenchError(f"Not enough instances available: {len(instances)} < {nodes}")
        
        try:
            # Build binary locally
            self._build_binary()
            
            # Create connections
            self._create_connections(instances[:nodes])
            
            # Install dependencies
            self._install_dependencies()
            
            # Upload code
            self._upload_code()
            
            # Build on remote instances
            self._build_remote()
            
            # Generate configurations
            self._generate_remote_configs(nodes, node_params)
            
            Print.success("Deployment completed successfully")
            
        except Exception as e:
            raise BenchError(f"Deployment failed: {e}")
    
    def _build_binary(self):
        """Build Dumbo binary locally"""
        Print.info("Building Dumbo binary locally...")
        
        returncode, stdout, stderr = run_command("go mod tidy", self.root_dir)
        if returncode != 0:
            raise BenchError(f"go mod tidy failed: {stderr}")
        
        returncode, stdout, stderr = run_command("go build -o dumbo .", self.root_dir)
        if returncode != 0:
            raise BenchError(f"go build failed: {stderr}")
        
        Print.success("Local binary built successfully")
    
    def _create_connections(self, instances: List[Dict[str, Any]]):
        """Create SSH connections to instances"""
        Print.info("Creating SSH connections...")
        
        self.connections = []
        for i, instance in enumerate(instances):
            try:
                conn = Connection(
                    host=instance['public_ip'],
                    user='root',  # Alibaba Cloud typically uses root
                    connect_kwargs={'key_filename': 'dumbo-alibaba-key.pem'}
                )
                self.connections.append(conn)
                Print.info(f"Connected to instance {i}: {instance['public_ip']}")
            except Exception as e:
                raise BenchError(f"Failed to connect to instance {i}: {e}")
    
    def _install_dependencies(self):
        """Install dependencies on all instances"""
        Print.info("Installing dependencies on Alibaba Cloud instances...")
        
        for i, conn in enumerate(self.connections):
            try:
                # Update system
                conn.run('apt-get update', hide=True)
                
                # Install Go
                conn.run('apt-get install -y golang-go', hide=True)
                
                # Install tmux
                conn.run('apt-get install -y tmux', hide=True)
                
                # Install Python3 and pip
                conn.run('apt-get install -y python3 python3-pip', hide=True)
                
                Print.info(f"Dependencies installed on instance {i}")
                
            except Exception as e:
                raise BenchError(f"Failed to install dependencies on instance {i}: {e}")
    
    def _upload_code(self):
        """Upload Dumbo code to instances"""
        Print.info("Uploading code to Alibaba Cloud instances...")
        
        # Create tar archive
        returncode, stdout, stderr = run_command("tar -czf dumbo.tar.gz --exclude='data' --exclude='*.log' .", self.root_dir)
        if returncode != 0:
            raise BenchError(f"Failed to create tar archive: {stderr}")
        
        for i, conn in enumerate(self.connections):
            try:
                # Upload tar file
                conn.put('dumbo.tar.gz', '/root/dumbo.tar.gz')
                
                # Extract and setup
                conn.run('cd /root && tar -xzf dumbo.tar.gz && rm dumbo.tar.gz', hide=True)
                conn.run('cd /root && mkdir -p data', hide=True)
                
                Print.info(f"Code uploaded to instance {i}")
                
            except Exception as e:
                raise BenchError(f"Failed to upload code to instance {i}: {e}")
    
    def _build_remote(self):
        """Build Dumbo on remote instances"""
        Print.info("Building Dumbo on remote instances...")
        
        for i, conn in enumerate(self.connections):
            try:
                conn.run('cd /root && go mod tidy', hide=True)
                conn.run('cd /root && go build -o dumbo .', hide=True)
                
                Print.info(f"Dumbo built on instance {i}")
                
            except Exception as e:
                raise BenchError(f"Failed to build Dumbo on instance {i}: {e}")
    
    def _generate_remote_configs(self, nodes: int, node_params: Dict[str, Any]):
        """Generate configurations on remote instances"""
        Print.info("Generating configurations on remote instances...")
        
        for i, conn in enumerate(self.connections):
            try:
                # Generate config template
                config_template = f"""
id_name:
{chr(10).join([f"  {j}: node{j}" for j in range(nodes)])}
id_ip:
{chr(10).join([f"  {j}: {self.instance_manager.instances[j]['private_ip']}" for j in range(nodes)])}
id_p2p_port:
{chr(10).join([f"  {j}: {9000 + j}" for j in range(nodes)])}
log_level: {node_params['dumbo']['log_level']}
max_pool: 10
timeout: 1000
mock_latency: {node_params['dumbo']['mock_latency']}
ddos: true
max_payload_size: {node_params['dumbo']['max_payload_size']}
tx_size: {node_params['dumbo']['tx_size']}
wait_time: {node_params['dumbo']['wait_time']}
"""
                
                # Write config template
                conn.run(f'cd /root && cat > config_temp.yaml << "EOF"\n{config_template}\nEOF', hide=True)
                
                # Generate individual configs
                conn.run('cd /root/config_gen && cp ../config_temp.yaml ./config_temp.yaml', hide=True)
                conn.run('cd /root/config_gen && go run main.go', hide=True)
                conn.run('cd /root/config_gen && sed -i "s/TSPubKey:/tspubkey:/g; s/TSShare:/tsshare:/g" *.yaml', hide=True)
                conn.run('cd /root/config_gen && cp *.yaml ..', hide=True)
                
                Print.info(f"Configurations generated on instance {i}")
                
            except Exception as e:
                raise BenchError(f"Failed to generate configurations on instance {i}: {e}")
    
    def _run_benchmark(self, bench_params: Dict[str, Any], node_params: Dict[str, Any]):
        """Run benchmark on Alibaba Cloud instances"""
        Print.subheading("Running benchmark on Alibaba Cloud instances")
        
        nodes = bench_params['nodes'][0]
        duration = bench_params['duration']
        
        try:
            # Start nodes
            for i, conn in enumerate(self.connections):
                if i < nodes:
                    session = f"dumbo_{i}"
                    cmd = f'cd /root && ./dumbo -config node{i}'
                    conn.run(f'tmux new -d -s "{session}" "{cmd} > data/logs/node-{i}.log 2>&1"', hide=True)
                    Print.info(f"Started node {i} on instance {i}")
            
            # Wait for benchmark
            Print.info(f"Running benchmark for {duration} seconds...")
            time.sleep(duration)
            
            # Stop nodes
            for i, conn in enumerate(self.connections):
                conn.run('tmux kill-server || true', hide=True)
                Print.info(f"Stopped node {i} on instance {i}")
            
            Print.success("Benchmark execution completed")
            
        except Exception as e:
            raise BenchError(f"Benchmark execution failed: {e}")
    
    def _collect_results(self):
        """Collect results from Alibaba Cloud instances"""
        Print.subheading("Collecting results from Alibaba Cloud instances")
        
        timestamp = get_timestamp()
        results_dir = os.path.join(self.data_dir, f"alibaba_{timestamp}")
        ensure_directory(results_dir)
        ensure_directory(os.path.join(results_dir, "logs"))
        
        try:
            for i, conn in enumerate(self.connections):
                # Download logs
                conn.get('/root/data/logs/', f'{results_dir}/logs/', preserve_mode=True)
                Print.info(f"Downloaded logs from instance {i}")
            
            # Parse results
            self._parse_alibaba_results(results_dir, len(self.connections))
            
            Print.success(f"Results collected to: {results_dir}")
            
        except Exception as e:
            raise BenchError(f"Failed to collect results: {e}")
    
    def _parse_alibaba_results(self, results_dir: str, nodes: int):
        """Parse Alibaba Cloud benchmark results"""
        Print.info("Parsing Alibaba Cloud benchmark results...")
        
        # Similar to local parsing but adapted for Alibaba Cloud results
        logs_dir = os.path.join(results_dir, "logs")
        
        if not os.path.exists(logs_dir):
            Print.warning("No logs directory found")
            return
        
        # Parse logs and generate results
        # Implementation similar to local benchmark parsing
        Print.success("Alibaba Cloud results parsed successfully")
    
    def _cleanup(self):
        """Cleanup Alibaba Cloud resources"""
        Print.info("Cleaning up Alibaba Cloud resources...")
        
        try:
            # Close connections
            for conn in self.connections:
                conn.close()
            self.connections = []
            
            # Terminate instances
            self.instance_manager.terminate_instances()
            
            Print.success("Cleanup completed")
            
        except Exception as e:
            Print.warning(f"Cleanup failed: {e}")

